import psutil
import prometheus_client
from prometheus_client import Gauge
from flask import Response, Flask

# Initialize metrics
cpu_usage = Gauge("system_cpu_usage", "CPU usage percentage")
memory_usage = Gauge("system_memory_usage", "Memory usage percentage")
disk_usage = Gauge("system_disk_usage", "Disk usage percentage", ['device'])
load_average = Gauge("system_load_average", "System load average", ['interval'])
network_bytes = Gauge("system_network_bytes", "Network traffic in bytes", ['interface', 'direction'])

app = Flask(__name__)


def collect_metrics():
    # CPU usage
    cpu_usage.set(psutil.cpu_percent(interval=1))

    # Memory usage
    memory = psutil.virtual_memory()
    memory_usage.set(memory.percent)

    # Disk usage
    for partition in psutil.disk_partitions():
        usage = psutil.disk_usage(partition.mountpoint)
        disk_usage.labels(device=partition.device).set(usage.percent)

    # Load average
    load_avg = psutil.getloadavg()
    load_average.labels(interval='1m').set(load_avg[0])
    load_average.labels(interval='5m').set(load_avg[1])
    load_average.labels(interval='15m').set(load_avg[2])

    # Network traffic
    net_io = psutil.net_io_counters(pernic=True)
    for interface, counters in net_io.items():
        network_bytes.labels(interface=interface, direction='sent').set(counters.bytes_sent)
        network_bytes.labels(interface=interface, direction='received').set(counters.bytes_recv)


@app.route("/metrics")
def requests_metrics():
    collect_metrics()
    return Response(prometheus_client.generate_latest(), mimetype="text/plain")


@app.route('/')
def index():
    html = '''<html>
            <head><title>System Metrics Exporter</title></head>
            <body>
            <h1>System Metrics Exporter</h1>
            <p><a href="/metrics">Metrics</a></p>
            </body>
            </html>'''
    return html


if __name__ == "__main__":
    app.run(
        host="0.0.0.0",
        port=9912,
        debug=True
    )
